package com.sjw.workflow.engine;

import com.sjw.workflow.core.AgentResult;
import com.sjw.workflow.core.WorkflowContext;
import com.sjw.workflow.dag.WorkflowDAG;
import com.sjw.workflow.dag.WorkflowNode;
// import org.slf4j.Logger;
// import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

/**
 * Workflow执行引擎
 * 负责执行DAG中定义的Agent工作流
 */
public class WorkflowEngine {
    
    // private static final Logger logger = LoggerFactory.getLogger(WorkflowEngine.class);
    
    private final ExecutorService executorService;
    private final AtomicBoolean isRunning;
    
    public WorkflowEngine() {
        this.executorService = Executors.newCachedThreadPool();
        this.isRunning = new AtomicBoolean(false);
    }
    
    public WorkflowEngine(ExecutorService executorService) {
        this.executorService = executorService;
        this.isRunning = new AtomicBoolean(false);
    }
    
    /**
     * 执行Workflow
     */
    public CompletableFuture<WorkflowExecutionResult> execute(WorkflowDAG dag, WorkflowContext context) {
        return CompletableFuture.supplyAsync(() -> {
            if (!isRunning.compareAndSet(false, true)) {
                throw new IllegalStateException("Workflow engine is already running");
            }
            
            try {
                System.out.println("Starting workflow execution: " + dag.getWorkflowId());
                context.addLog("Workflow execution started");
                
                // 验证DAG
                if (!dag.validate()) {
                    throw new IllegalArgumentException("Invalid DAG structure: contains cycles");
                }
                
                // 执行工作流
                WorkflowExecutionResult result = executeWorkflow(dag, context);
                
                System.out.println("Workflow execution completed: " + dag.getWorkflowId());
                context.addLog("Workflow execution completed");
                
                return result;
                
            } finally {
                isRunning.set(false);
            }
        }, executorService);
    }
    
    /**
     * 执行工作流的核心逻辑
     */
    private WorkflowExecutionResult executeWorkflow(WorkflowDAG dag, WorkflowContext context) {
        long startTime = System.currentTimeMillis();
        int totalNodes = dag.size();
        int completedNodes = 0;
        int failedNodes = 0;
        
        try {
            while (completedNodes + failedNodes < totalNodes) {
                // 获取可以执行的节点
                List<WorkflowNode> readyNodes = dag.getReadyNodes(context);
                
                if (readyNodes.isEmpty()) {
                    // 检查是否所有节点都已完成
                    if (completedNodes + failedNodes >= totalNodes) {
                        break;
                    }
                    
                    // 如果没有可执行的节点，但还有未完成的节点，说明存在死锁
                    throw new RuntimeException("Workflow deadlock detected: no nodes can execute");
                }
                
                // 并行执行所有就绪的节点
                List<CompletableFuture<AgentResult>> futures = readyNodes.stream()
                        .map(node -> executeNode(node, context))
                        .collect(Collectors.toList());
                
                // 等待所有节点完成
                CompletableFuture<Void> allFutures = CompletableFuture.allOf(
                        futures.toArray(new CompletableFuture[0]));
                
                try {
                    allFutures.get(30, TimeUnit.SECONDS); // 30秒超时
                } catch (Exception e) {
                    System.err.println("Node execution timeout or error: " + e.getMessage());
                    throw new RuntimeException("Node execution failed", e);
                }
                
                // 统计结果
                for (CompletableFuture<AgentResult> future : futures) {
                    try {
                        AgentResult result = future.get();
                        if (result.getStatus() == com.sjw.workflow.core.ExecutionStatus.SUCCESS) {
                            completedNodes++;
                        } else {
                            failedNodes++;
                        }
                    } catch (Exception e) {
                        System.err.println("Failed to get node result: " + e.getMessage());
                        failedNodes++;
                    }
                }
                
                System.out.println("Progress: " + completedNodes + "/" + totalNodes + " nodes completed, " + failedNodes + " failed");
            }
            
            long executionTime = System.currentTimeMillis() - startTime;
            
            return new WorkflowExecutionResult(
                    dag.getWorkflowId(),
                    completedNodes,
                    failedNodes,
                    totalNodes,
                    executionTime,
                    context
            );
            
        } catch (Exception e) {
            long executionTime = System.currentTimeMillis() - startTime;
            System.err.println("Workflow execution failed: " + e.getMessage());
            
            return new WorkflowExecutionResult(
                    dag.getWorkflowId(),
                    completedNodes,
                    failedNodes,
                    totalNodes,
                    executionTime,
                    context,
                    e.getMessage()
            );
        }
    }
    
    /**
     * 执行单个节点
     */
    private CompletableFuture<AgentResult> executeNode(WorkflowNode node, WorkflowContext context) {
        return CompletableFuture.supplyAsync(() -> {
            System.out.println("Executing node: " + node.getId() + " (" + node.getName() + ")");
            context.addLog("Executing node: " + node.getId());
            
            try {
                // 获取节点输入
                com.sjw.workflow.core.AgentInput input = node.getInput(context);
                
                // 执行Agent
                CompletableFuture<AgentResult> future = node.getAgent().execute(input, context);
                AgentResult result = future.get(60, TimeUnit.SECONDS); // 60秒超时
                
                // 保存结果到上下文
                context.setAgentResult(node.getId(), result);
                
                System.out.println("Node " + node.getId() + " completed with status: " + result.getStatus());
                context.addLog("Node " + node.getId() + " completed with status: " + result.getStatus());
                
                return result;
                
            } catch (Exception e) {
                System.err.println("Node " + node.getId() + " execution failed: " + e.getMessage());
                context.addLog("Node " + node.getId() + " execution failed: " + e.getMessage());
                
                AgentResult errorResult = new AgentResult(node.getId(), 
                        com.sjw.workflow.core.ExecutionStatus.FAILED, 
                        new java.util.HashMap<>());
                errorResult.setError(e.getMessage());
                context.setAgentResult(node.getId(), errorResult);
                
                return errorResult;
            }
        }, executorService);
    }
    
    /**
     * 停止执行引擎
     */
    public void shutdown() {

        executorService.shutdown();
        try {
            if (!executorService.awaitTermination(30, TimeUnit.SECONDS)) {
                executorService.shutdownNow();
            }
        } catch (InterruptedException e) {
            executorService.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }
    
    /**
     * 检查引擎是否正在运行
     */
    public boolean isRunning() {
        return isRunning.get();
    }
}

/**
 * Workflow执行结果
 */
